import shutil
from openmind_hub import create_repo, delete_repo, http_get, om_hub_url, OM_HUB_CACHE, om_hub_download, \
    try_to_load_from_cache, snapshot_download,upload_file,upload_folder
from openmind_hub import build_om_headers
import os,os.path
import string,random
import unittest
import time
import tempfile
from functools import partial
from openmind_hub import RepositoryNotFoundError, EntryNotFoundError, GatedRepoError

token = '123456789()'
repo_id = "modelfoundryinfra/Qwen-14B-Chat"
def search_tmp_File(pathname):
    matchedFile =[]
    for root,dirs,files in os.walk(pathname):
        for file in files:
            if file.startswith("tmp"):
                fname = os.path.abspath(os.path.join(root, file))
                matchedFile.append(fname)
        return matchedFile
def search_dir_File(pathname):
    folders = [f for f in os.listdir(pathname) if os.path.isdir(os.path.join(pathname, f))]
    filtered_folders = [f for f in folders if 'openmind' in f or 'model' in f or 'space' in f or 'dataset' in f]
    return filtered_folders
def check_file_exist(path):
    return os.path.exists(path)
def generate_large_file(filename, size_in_bytes):
    with open(filename, 'wb') as file:
        chunk_size = 1024 * 1024  # 1MB
        chunks = size_in_bytes // chunk_size
        remainder = size_in_bytes % chunk_size
        data = b'\0' * chunk_size

        for _ in range(chunks):
            file.write(data)

        if remainder > 0:
            file.write(data[:remainder])

class TestOpenmindHub(unittest.TestCase):
    def test_01_httpget_public_success(self):
        try:
            temp_file_manager = partial(
                tempfile.NamedTemporaryFile, mode="wb", dir=OM_HUB_CACHE, delete=False
            )
            with temp_file_manager() as temp_file:
                http_get(
                    url=om_hub_url(repo_id=repo_id, filename="config.json"),
                    temp_file=temp_file,
                    target_dir=OM_HUB_CACHE,
                    proxies=None,
                    resume_size=0,
                    headers=None,
                    expected_size=None,
                    _nb_retries=5,
                )
            time.sleep(3)
            files=search_tmp_File(OM_HUB_CACHE)
            self.assertTrue(len(files),"下载的文件不存在")
        except:
           raise Exception("下载文件有报错")
        else:
            for file in files:
                os.remove(file)
    def test_02_httpget_public_noexistrepo(self):
        try:
            temp_file_manager = partial(
                tempfile.NamedTemporaryFile, mode="wb", dir=OM_HUB_CACHE, delete=False
            )
            with temp_file_manager() as temp_file:
                http_get(
                    url=om_hub_url(repo_id=repo_id+"a", filename="config.json"),
                    temp_file=temp_file,
                    target_dir=OM_HUB_CACHE,
                    proxies=None,
                    resume_size=0,
                    headers=None,
                    expected_size=None,
                    _nb_retries=5,
                )
            time.sleep(3)
            files=search_tmp_File(OM_HUB_CACHE)
            self.assertFalse(len(files),"下载的文件不应该存在")
        except RepositoryNotFoundError:
           pass
    def test_03_httpget_public_noexistfile(self):
        try:
            temp_file_manager = partial(
                tempfile.NamedTemporaryFile, mode="wb", dir=OM_HUB_CACHE, delete=False
            )
            with temp_file_manager() as temp_file:
                http_get(
                    url=om_hub_url(repo_id=repo_id, filename="invalid.json"),
                    temp_file=temp_file,
                    target_dir=OM_HUB_CACHE,
                    proxies=None,
                    resume_size=0,
                    headers=None,
                    expected_size=None,
                    _nb_retries=5,
                )
        except EntryNotFoundError:
           pass
    def test_04_httpget_private_success(self):
        repo_id = "guoxiaozhen2/model_test2024"
        license = "mit"
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            temp_file_manager = partial(
                tempfile.NamedTemporaryFile, mode="wb", dir=OM_HUB_CACHE, delete=False
            )
            with temp_file_manager() as temp_file:
                http_get(
                    url=om_hub_url(repo_id=repo_id, filename="README.md"),
                    temp_file=temp_file,
                    target_dir=OM_HUB_CACHE,
                    proxies=None,
                    resume_size=0,
                    headers=build_om_headers(token=token),
                    expected_size=None,
                    _nb_retries=5,
                )
            time.sleep(3)
            files=search_tmp_File(OM_HUB_CACHE)
            self.assertTrue(len(files),"下载的文件不存在")
            for file in files:
                os.remove(file)
        except:
           raise Exception("下载文件有报错")
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
    def test_05_httpget_private_notoken(self):
        repo_id = "guoxiaozhen2/model_test2024"
        license = "mit"
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            temp_file_manager = partial(
                tempfile.NamedTemporaryFile, mode="wb", dir=OM_HUB_CACHE, delete=False
            )
            with temp_file_manager() as temp_file:
                http_get(
                    url=om_hub_url(repo_id=repo_id, filename="README.md"),
                    temp_file=temp_file,
                    target_dir=OM_HUB_CACHE,
                    proxies=None,
                    resume_size=0,
                    headers=None,
                    expected_size=None,
                    _nb_retries=5,
                )
        except RepositoryNotFoundError:
           pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
    def test_06_om_hub_download_success(self):
        try:
            om_hub_download(
                repo_id=repo_id,
                filename="README.md",
                subfolder=None,
                user_agent=None,
                force_download=False,
                token=token,
                local_files_only=False,
            )
            time.sleep(3)
            dirs=search_dir_File(OM_HUB_CACHE)
            self.assertTrue(len(dirs),"下载的文件不存在")
        except:
           raise Exception("下载文件有报错")
        else:
            for folder in dirs:
                folder_to_delete = os.path.join(OM_HUB_CACHE, folder)
                shutil.rmtree(folder_to_delete)
    def test_07_om_hub_download_wrongtoken(self):
        try:
            om_hub_download(
                repo_id=repo_id,
                filename="README.md",
                subfolder=None,
                user_agent=None,
                force_download=False,
                token=token+"A",
                local_files_only=False,
            )
            time.sleep(3)
            dirs=search_dir_File(OM_HUB_CACHE)
            self.assertFalse(len(dirs),"下载的文件不存在")
        except EntryNotFoundError:
           pass
    def test_08_om_hub_download_noexistrepo(self):
        try:
            om_hub_download(
                repo_id=repo_id+"A",
                filename="README.md",
                subfolder=None,
                user_agent=None,
                force_download=False,
                token=token,
                local_files_only=False,
            )
            time.sleep(3)
            dirs=search_dir_File(OM_HUB_CACHE)
            self.assertFalse(len(dirs),"下载的文件不存在")
        except RepositoryNotFoundError:
           pass
    def test_09_try_to_load_from_cache_success(self):
        try:
            om_hub_download(
                repo_id=repo_id,
                filename="README.md",
                subfolder=None,
                user_agent=None,
                force_download=False,
                token=token,
                local_files_only=False,
            )
            time.sleep(3)
            dirs = search_dir_File(OM_HUB_CACHE)
            a = try_to_load_from_cache(
                repo_id=repo_id,
                filename="README.md",
            )
            self.assertTrue(a, "未正确返回文件路径")
        except:
           raise Exception("查找文件在缓存中的路径有报错")
        else:
            for folder in dirs:
                folder_to_delete = os.path.join(OM_HUB_CACHE, folder)
                shutil.rmtree(folder_to_delete)
    def test_10_try_to_load_from_cache_noexistrepo(self):
        try:
            om_hub_download(
                repo_id=repo_id,
                filename="README.md",
                subfolder=None,
                user_agent=None,
                force_download=False,
                token=token,
                local_files_only=False,
            )
            time.sleep(3)
            dirs = search_dir_File(OM_HUB_CACHE)
            a = try_to_load_from_cache(
                repo_id=repo_id+"a",
                filename="README.md",
            )
            self.assertFalse(a, "未正确返回文件路径")
        except:
           raise Exception("查找文件在缓存中的路径有报错")
        else:
            for folder in dirs:
                folder_to_delete = os.path.join(OM_HUB_CACHE, folder)
                shutil.rmtree(folder_to_delete)
    def test_11_try_to_load_from_cache_noexistfile(self):
        try:
            om_hub_download(
                repo_id=repo_id,
                filename="README.md",
                subfolder=None,
                user_agent=None,
                force_download=False,
                token=token,
                local_files_only=False,
            )
            time.sleep(3)
            dirs = search_dir_File(OM_HUB_CACHE)
            a = try_to_load_from_cache(
                repo_id=repo_id,
                filename="invalid.txt",
            )
            self.assertFalse(a, "未正确返回文件路径")
        except:
           raise Exception("查找文件在缓存中的路径有报错")
        else:
            for folder in dirs:
                folder_to_delete = os.path.join(OM_HUB_CACHE, folder)
                shutil.rmtree(folder_to_delete)
    def test_12_snapshot_download_cachedir_success(self):
        repo_id = "guoxiaozhen2/model_test2024"
        license = "mit"
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            snapshot_download(
                repo_id=repo_id,
                repo_type=None,
                revision=None,
                cache_dir=OM_HUB_CACHE,
                #local_dir="space_test0516",
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            time.sleep(3)
            dirs=search_dir_File(OM_HUB_CACHE)
            self.assertTrue(len(dirs),"下载的文件不存在")
        except:
           raise Exception("下载文件有报错")
        else:
            for folder in dirs:
                folder_to_delete = os.path.join(OM_HUB_CACHE, folder)
                shutil.rmtree(folder_to_delete)
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
    def test_13_snapshot_download_wrong_token(self):
        repo_id = "guoxiaozhen2/model_test2024"
        license = "mit"
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            snapshot_download(
                repo_id=repo_id,
                repo_type=None,
                revision=None,
                cache_dir=OM_HUB_CACHE,
                #local_dir="space_test0516",
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token+"a",
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            raise Exception("下载文件有报错")
        except GatedRepoError:
           pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
    def test_14_snapshot_download_noexist_repo(self):
        repo_id = "guoxiaozhen2/model_test2024"
        license = "mit"
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            snapshot_download(
                repo_id=repo_id+"a",
                repo_type=None,
                revision=None,
                cache_dir=OM_HUB_CACHE,
                #local_dir="space_test0516",
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            raise Exception("下载文件有报错")
        except RepositoryNotFoundError:
           pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
    def test_15_snapshot_download_localdir_success(self):
        repo_id = "guoxiaozhen2/model_test2024"
        author = "guoxiaozhen2"
        license = "mit"
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            snapshot_download(
                repo_id=repo_id,
                repo_type=None,
                revision=None,
                # cache_dir=OM_HUB_CACHE,
                local_dir=r"/openmind-hub/"+repo_id,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            time.sleep(3)
            dirs=search_dir_File(r"/openmind-hub/"+author)
            self.assertTrue(len(dirs),"下载的文件不存在")
        except:
           raise Exception("下载文件有报错")
        else:
            shutil.rmtree(r"/openmind-hub/"+repo_id)
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
    def test_16_snapshot_download_type_id_notconsistent(self):
        repo_id = "guoxiaozhen2/model_test2024"
        license = "mit"
        try:
            create_repo(
                repo_id=repo_id,
                private=True,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            snapshot_download(
                repo_id=repo_id,
                repo_type="dataset",
                revision=None,
                # cache_dir=OM_HUB_CACHE,
                local_dir=r"D:/hub/"+repo_id,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            raise Exception("没有校验type与ID不一致的情况")
        except RepositoryNotFoundError:
           pass
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
    def test_17_snapshot_download_localdir_largefile_5G(self):
        repo_id = "guoxiaozhen2/model_test2024"
        author = "guoxiaozhen2"
        license = "mit"
        file_path = "/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/testfile"
        size_in_bytes = int(5 * 1024 * 1024 * 1024)
        try:
            generate_large_file(file_path, size_in_bytes)
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=file_path,
                path_in_repo="a/filedata.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            snapshot_download(
                repo_id=repo_id,
                repo_type=None,
                revision=None,
                # cache_dir=OM_HUB_CACHE,
                local_dir=r"/openmind-hub/"+repo_id,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            time.sleep(3)
            dirs=search_dir_File(r"/openmind-hub/"+author)
            self.assertTrue(len(dirs),"下载的文件不存在")
        except:
           raise Exception("下载文件有报错")
        else:
            shutil.rmtree(r"/openmind-hub/"+author)
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
            os.remove(file_path)
    def test_18_snapshot_download_localdir_largefile_50G(self):
        repo_id = "guoxiaozhen2/model_test2024"
        author = "guoxiaozhen2"
        license = "mit"
        file_path = "/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/testfile"
        size_in_bytes = int(50 * 1024 * 1024 * 1024)
        try:
            generate_large_file(file_path, size_in_bytes)
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="model",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_file(
                path_or_fileobj=file_path,
                path_in_repo="a/filedata.txt",
                repo_id=repo_id,
                token=token,
                revision="main",
                commit_message="upload file",
                commit_description=None,
            )
            time.sleep(3)
            snapshot_download(
                repo_id=repo_id,
                repo_type=None,
                revision=None,
                # cache_dir=OM_HUB_CACHE,
                local_dir=r"/openmind-hub/"+repo_id,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            time.sleep(3)
            dirs=search_dir_File(r"/openmind-hub/"+author)
            self.assertTrue(len(dirs),"下载的文件不存在")
        except:
           raise Exception("下载文件有报错")
        else:
            shutil.rmtree(r"/openmind-hub/"+author)
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="model",
                missing_ok=True
            )
            os.remove(file_path)

    def test_19_snapshotdownload_multifolder_2000(self):
        repo_id = "guoxiaozhen2/model_test2024"
        author = "guoxiaozhen2"
        license = "mit"
        file_dir = []
        # 创建2000个随机文件
        for i in range(2000):
            # 生成随机文件名
            random_file_name = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
            file_path = f"/openmind-hub/auto_test/autotest/test_openmind_hub/filedata/{random_file_name}.txt"
            file_dir.append(file_path)
            # 创建空白文件
            with open(file_path, 'w'):
                pass
        print("已创建2000个随机文件")
        try:
            create_repo(
                repo_id=repo_id,
                private=False,
                desc="描述内容",
                license=license,
                token=token,
                repo_type="dataset",
                exist_ok=True,
                fullname=None
            )
            time.sleep(3)
            upload_folder(
                repo_id=repo_id,
                folder_path="/openmind-hub/auto_test/autotest/test_openmind_hub/filedata",
                path_in_repo="/filedata",
                commit_message="upload folder",
                commit_description=None,
                token=token,
                revision=None,
                # allow_patterns=["*.bin", "*.py"],
                ignore_patterns="*.log",
                num_threads=5,
            )
            time.sleep(3)
            snapshot_download(
                repo_id=repo_id,
                repo_type="dataset",
                revision=None,
                # cache_dir=OM_HUB_CACHE,
                local_dir=r"/openmind-hub/auto_test/autotest/" + repo_id,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            time.sleep(3)
            dirs = search_dir_File(r"/openmind-hub/auto_test/autotest/" + author)
            self.assertTrue(len(dirs), "下载的文件不存在")
        except:
            raise Exception("下载文件有报错")
        # else:
        #     shutil.rmtree(r"./" + author)
        finally:
            delete_repo(
                repo_id=repo_id,
                token=token,
                repo_type="dataset",
                missing_ok=True
            )
            for i in file_dir:
                os.remove(i)

if __name__ == '__main__':
    unittest.main()